// (C) 2022 Intel Corporation. All rights reserved.
// Your use of Intel Corporation's design tools, logic functions and other 
// software and tools, and its AMPP partner logic functions, and any output 
// files from any of the foregoing (including device programming or simulation 
// files), and any associated documentation or information are expressly subject 
// to the terms and conditions of the Intel Program License Subscription 
// Agreement, Intel FPGA IP License Agreement, or other applicable 
// license agreement, including, without limitation, that your use is for the 
// sole purpose of programming logic devices manufactured by Intel and sold by 
// Intel or its authorized distributors.  Please refer to the applicable 
// agreement for further details.

module A_cse_ocs_dpa_aes10c_control #(
   //Added parameter for AES', for now only used for DVB-CSA3
   parameter     AES_P = 1'b0
)
(
   input  logic         sys_clk,              // system clock
   input  logic         reset_n,              // system reset
   output logic         aes_start,            // Start AES operation
   input  logic [1:0]   aes_op,               // AES Operation: 0 - Encrypt, 1 - Decrypt, 2 - Expand, 3 - Bypass 
   input  logic		DMA_Sync_Reset,
   input  logic [1:0]   i_key_size,
   input  logic         data_valid,
   output logic         aes_busy,             // AES busy indicator
   output logic         encrypt,              // Control Signal to state module to enable encrypt data-path
   output logic         mixcolumns,           // do mixcolumns operation (or inverse)
   output logic         bytesub_shiftrows,    // do byte substitution and row shifting (or inverse)
   output logic         add_roundkey,         // Added for AES', always 1 for AES
   output logic         store_state,          // store the state
   output logic         store_key,            // store the key
   output logic         sbox_input_sel,       // key register input selector for sbox
   output logic [7:0]   rcon,                 // upper 8 bits of the round constant
   output logic         even_rounds,   
   output logic         odd_rounds,
   output logic         lastround,
   output logic         aes_key_start,
   output logic         load_expanded
   );

`include "A_cse_ocs_dpa_aes10c_localparams.v"   
/**********************************************************************/
// Port Declarations
/**********************************************************************/
  
   parameter  X     = 1'bz;
   parameter  H     = 1'b1;
   parameter  L     = 1'b0;

/**********************************************************************/
// Data Type Declarations (Signal Declarations)
/**********************************************************************/
  reg 		 aes_start_r;
  reg [3:0] 	 round;                       // round counter
  reg [3:0] 	 keyround;                    // keyround counter
  reg 		 expansion;                   // indicates expansion operation in progress
  reg 		 extra_round;                 // Only set if AES_P is set
  reg		 add_key;
  
  wire [3:0] 	 rcon_key_round;
  wire 		 decrypt;                     // 1 for decryption
  logic 	 firstround;                  // indicates 1st round of a cryption operation

/**********************************************************************/
// Main code
/**********************************************************************/

assign decrypt = !encrypt;

// hardware registers
always_ff @ (posedge sys_clk or negedge reset_n) 
    if (~reset_n) 
      aes_start <= 1'b0;
    else if(DMA_Sync_Reset)
      aes_start <= 1'b0;
    else if((data_valid & lastround) | (data_valid & ~aes_busy)) 
      aes_start <= 1'b1;
    else if(aes_start)
      aes_start <= 1'b0;
 
 
always_ff @ (posedge sys_clk or negedge reset_n) 
    if (~reset_n) 
      aes_start_r  <= 1'b0;
    else if(DMA_Sync_Reset)
      aes_start_r  <= 1'b0;
    else 
      aes_start_r <= (data_valid & ~aes_busy);

assign aes_key_start = (data_valid & lastround) | ((data_valid & ~aes_busy) & ~aes_start_r);     

always_ff @ (posedge sys_clk or negedge reset_n) begin: internal_registers
    if (~reset_n) begin
        aes_busy          <= 1'b0;
        round             <= 4'h0;
        keyround          <= 4'h0;
        expansion         <= 1'b0;
        encrypt           <= 1'b0;
        extra_round       <= 1'b0;
    end 
    else if(DMA_Sync_Reset) begin
        aes_busy          <= 1'b0;
        round             <= 4'h0;
        keyround          <= 4'h0;
        expansion         <= 1'b0;
        encrypt           <= 1'b0;
        extra_round       <= 1'b0;
    end
	else begin
        // aes_busy flag -
        // Set when encrypt/decrypt/expand op is started
        // Clear when crypt operation is done.
        if (lastround & ~data_valid)
          aes_busy <= 1'b0;
        else if (data_valid & ~aes_busy) 
          aes_busy <= 1'b1;

        // Check if an expansion or encryption operation is started
        if (aes_start & (aes_op == 2'b10))
          encrypt <= 1'b1; //Key Expansion
        else if (data_valid & (aes_op == 2'b00)) 
          encrypt <=  AES_P ? 1'b0 : 1'b1; //Encrypt for AES, decrypt for AES'
        else if (data_valid & (aes_op == 2'b01))
          encrypt <=  AES_P ? 1'b1 : 1'b0; //Decrypt for AES, encrypt for AES'

        // Check if an expansion operation is started
        if (aes_start & aes_op[1])
          expansion <= 1'b1;
        else if (aes_start)
          expansion <= 1'b0;

        // State round counter -
        // Keeps track of what cryption round we're in. Set to round 0 when we start
        // (round 0 actually corresponds to the 1st addroundkey operation)
        if (lastround)
          round <= 4'h0;
        else if (aes_busy) begin
            if (lastround)
              round <= 4'h0;
            else if (AES_P && (round == 4'h5))
              round <= round + {3'b0, extra_round};
            else 
              round <= round + 4'h1;
        end

        //Extra round for AES'
        if (aes_start)
           extra_round  <= 1'b0;
        else 
           extra_round  <= (round == 4'h5) ? !extra_round  : extra_round ;

        // Key round counter - 
        // Keeps track of the keyround which is different from the state
        // processing "round".  The keyround determines the round constant
        // value that gets used during key expansion/de-expansion so it
        // gets incremented differently depending on key size.
        // 128-bit key => increment at the end of every state round
        // 256-bit key => increment at the end of every other state round
        if (lastround)
          keyround <= 4'h0;
        // Added stalling keyround for the extra round in AESP mode
        else if (aes_busy & !(AES_P & (round==4'h5) & !extra_round))
          keyround <= keyround + 4'h1;
   end
end

// first/last round indicators
assign firstround = (round == 4'h0);

assign lastround = (aes_op == 2'b11) ? (firstround) : ((i_key_size == NK_128) && (round == NR_128)) ||
                   ((i_key_size == NK_256) && (round == NR_256));

assign load_expanded = ((i_key_size == NK_128) && (round == (NR_128-4'h1))) || ((i_key_size == NK_256) && (round == (NR_256-4'h1)));

// even/odd round indicators for 256-bit key processing
assign even_rounds = ~round[0];
assign odd_rounds  = ~even_rounds;

// Controls for processing the state
assign mixcolumns           = !expansion & !firstround & !lastround;
assign bytesub_shiftrows    = ((!expansion & (encrypt & !firstround)) | (decrypt & !lastround)) & aes_busy;
//Added for AES_P, no AddRoundKey for encryption extraround, and does AddRoundKey for decryption extraround
assign add_roundkey         = !AES_P ? add_key & aes_busy : encrypt ? !((round == 4'h5) && extra_round) : !((round == 4'h5) && !extra_round);

// Disable add_round key after last round for DPA resistance
always_ff @(posedge sys_clk or negedge reset_n)
  begin
    if (~reset_n)
	add_key <= 1'b0;
    else if(DMA_Sync_Reset)
	add_key <= 1'b0;
    else begin
       
       if (data_valid)
	  add_key <= 1'b1;
       else if (lastround)
          add_key <= 1'b0;
       else
	  add_key <= add_key;
    end
  end

// Store the state
assign store_state = !expansion & aes_busy;

assign sbox_input_sel = decrypt;

// RCON Calculation
assign rcon     = rcon_calc(encrypt,i_key_size,rcon_key_round);
assign rcon_key_round = (i_key_size == NK_128) ? keyround : {1'b0,keyround[3:1]};

// Store the key
assign store_key   = aes_busy & !lastround & !(AES_P & (round==4'h5) & !extra_round) & ((i_key_size == NK_128) | ((i_key_size == NK_256) & odd_rounds));

/**********************************************************************/
// Functions
/**********************************************************************/

/****************************************************************/
// Function name: rcon_calc
// Inputs       : encrypt  - 1 => encrypt, 0 => decrypt
//                key_size - 0 => 128, 1 => 192, 2 => 256
//                keyround - The key round. Range is from 0 to 9
//                           inclusive.
// Description  : Generates the upper 8 bits of the round constant 
//                (the lower 24 bits are 0).  Assumes that the
//                keyround always starts at 0 for both encrypt and
//                decrypt operations.
/****************************************************************/
function automatic [7:0] rcon_calc;
  input       encrypt;
  input [1:0] key_size;
  input [3:0] keyround;
  reg [3:0] rcon_addr,maxkeyround;
  begin
    case (key_size)
      NK_128:       maxkeyround = 4'h9;
      NK_192:       maxkeyround = 4'h7;
      NK_256,2'b11: maxkeyround = 4'h6;
      default : maxkeyround = 4'h9;
    endcase
    if (encrypt)
      rcon_addr = keyround;
    else
      rcon_addr = maxkeyround - keyround;

    //rcon_calc = 8'h00;
    case (rcon_addr) 
      4'h0: rcon_calc = 8'h01;
      4'h1: rcon_calc = 8'h02;
      4'h2: rcon_calc = 8'h04;
      4'h3: rcon_calc = 8'h08;
      4'h4: rcon_calc = 8'h10;
      4'h5: rcon_calc = 8'h20;
      4'h6: rcon_calc = 8'h40;
      4'h7: rcon_calc = 8'h80;
      4'h8: rcon_calc = 8'h1b;
      4'h9: rcon_calc = 8'h36;
      default : rcon_calc = 8'h00;
    endcase
  end
endfunction // rcon_calc

endmodule
